import torch
import torch.nn as nn
import torch.nn.functional as F


class FBNet(nn.Module):
    ''' FBNet '''
    def __init__(self):
        super(FBNet, self).__init__()
        # backend
        self.pre_layer = nn.Sequential(
            nn.Conv2d(3, 10, kernel_size=3, stride=1),  # conv1
            nn.PReLU(),  # PReLU1
            nn.MaxPool2d(kernel_size=2, stride=2),  # pool1
            nn.Conv2d(10, 16, kernel_size=3, stride=1),  # conv2
            nn.PReLU(),  # PReLU2
            nn.Conv2d(16, 32, kernel_size=3, stride=1),  # conv3
            nn.PReLU()  # PReLU3
        )
        # detection
        self.conv4_1 = nn.Conv2d(32, 1, kernel_size=1, stride=1)
        # bounding box regresion
        self.conv4_2 = nn.Conv2d(32, 4, kernel_size=1, stride=1)

        # weight initiation with xavier
        self.apply(weights_init)

    def forward(self, x, out_layer='conv6'):
        x = self.pre_layer(x)
        pred = F.softmax(self.conv4_1(x))
        # offset = self.conv4_2(x)

        return pred #, offset


def weights_init(m):
    if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
        nn.init.xavier_uniform_(m.weight.data)
        nn.init.constant_(m.bias, 0.1)


def fb_loss(gt_label, pred_label):
    # 定义损失函数
    loss_cls = nn.MSELoss()

    return loss_cls(gt_label, pred_label)